package bintree.avl;

import java.util.ArrayList;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * AVL树的实现
 * @author yuisama
 * @date 2022/03/07 18:33
 **/
public class AVLTree<E extends Comparable<E>> {
    // 节点类定义
    private class Node {
        private E val;
        private Node left;
        private Node right;
        // 当前结点的高度
        private int height;

        public Node(E val) {
            this.val = val;
            this.height = 1;
        }
    }
    private int size;
    // 根节点
    private Node root;

    public int getSize() {
        return size;
    }

    public void add(E newVal) {
        root = add(root,newVal);
    }

    /**
     * 当前节点的高度
     * @param node
     * @return
     */
    private int getHeight(Node node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 计算当前以node为根的BST的平衡因子
     * @param node
     * @return
     */
    private int getBalancedFactor(Node node) {
        if (node == null) {
            return 0;
        }
        return getHeight(node.left) - getHeight(node.right);
    }

    /**
     * 当前树是否是BST,BST的中序遍历结果是一个升序集合
     * @return
     */
    public boolean isBST() {
        List<E> list = new ArrayList<>();
        inOrder(root,list);
        for (int i = 0; i < list.size() - 1; i++) {
            if (list.get(i).compareTo(list.get(i + 1)) > 0) {
                System.err.println("not bst tree!");
                return false;
            }
        }
        return true;
    }

    public boolean isBalanced() {
        return isBalanced(root);
    }

    private boolean isBalanced(Node root) {
        if (root == null) {
            return true;
        }
        int balancedFactor = getBalancedFactor(root);
        if (Math.abs(balancedFactor) > 1) {
            return false;
        }
        return isBalanced(root.left) && isBalanced(root.right);
    }


    private void inOrder(Node root, List<E> list) {
        if (root == null) {
            return;
        }
        inOrder(root.left,list);
        list.add(root.val);
        inOrder(root.right,list);
    }


    /**
     * 向以root为根节点的BST中插入新元素,返回插入后的树根
     * @param root
     * @param newVal
     * @return 插入后的根节点
     */
    private Node add(Node root, E newVal) {
        if (root == null) {
            root = new Node(newVal);
            size ++;
        }
        if (newVal.compareTo(root.val) < 0) {
            // 左树插入
            root.left = add(root.left,newVal);
        }else if (newVal.compareTo(root.val) > 0){
            // 右树插入
            root.right = add(root.right,newVal);
        }
        // 更新当前节点的height
        root.height = 1 + Math.max(getHeight(root.left),getHeight(root.right));
        // 计算当前节点的平衡因子
        int balancedFactor = getBalancedFactor(root);
        // 平衡维护
        // LL
        if (balancedFactor > 1 && getBalancedFactor(root.left) >= 0) {
            // 右旋实现，此时插入元素后，不平衡的结点在当前树左侧的左侧,返回旋转后的根节点
            return rightRotate(root);
        }

        // RR
        if (balancedFactor < -1 && getBalancedFactor(root.right) <= 0) {
            // 左旋实现，此时插入元素后，不平衡的节点在当前树右侧的右侧，返回左旋后根节点
            return leftRotate(root);
        }

        // LR - 插入结点后，不平衡的结点在当前结点左树的右侧
        if (balancedFactor > 1 && getBalancedFactor(root.left) < 0) {
            // 先将左子树左旋转为LL
            root.left = leftRotate(root.left);
            // 再将当前root右旋
            return rightRotate(root);
        }

        // RL - 插入结点后，不平衡的结点在当前结点右树的左侧
        if (balancedFactor < -1 && getBalancedFactor(root.right) > 0) {
            // 先将右子树右旋转为RR
            root.right = rightRotate(root.right);
            // 再将当前root左旋
            return leftRotate(root);
        }
        // 当前结点满足平衡树的定义，继续回溯
        return root;
    }

    /**
     * 左旋实现
     * @param y
     * @return
     */
    private Node leftRotate(Node y) {
        Node x = y.right;
        Node T3 = x.left;
        x.left = y;
        y.right = T3;
        // 更新y和x的height
        y.height = 1 + Math.max(getHeight(y.left),getHeight(y.right));
        x.height = 1 + Math.max(getHeight(x.left),getHeight(x.right));
        return x;
    }

    /**
     * 右旋实现，返回右旋后的树根
     * @param y
     * @return
     */
    private Node rightRotate(Node y) {
        Node x = y.left;
        Node T3 = x.right;
        x.right = y;
        y.left = T3;
        // 更新x和y的height
        y.height = 1 + Math.max(getHeight(y.left),getHeight(y.right));
        x.height = 1 + Math.max(getHeight(x.left),getHeight(x.right));
        return x;
    }

    public boolean contains(E val) {
        return contains(root,val);
    }

    /**
     * 删除BST中值为e的节点
     * @param e
     */
    public void remove(E e) {
        root = remove(root,e);
    }

    /**
     * Hibbard Deletion
     * 在以root为根节点的BST中删除值为e的节点，返回删除后的树根
     * @param root
     * @param e
     * @return
     */
    private Node remove(Node root, E e) {
        if (root == null) {
            return null;
        }
        // 删除一个节点后，有可能破坏平衡，因此不能直接返回，需要最终检查平衡性
        Node retNode = null;
        if (e.compareTo(root.val) < 0) {
            // 左树删除
            root.left = remove(root.left,e);
            retNode = root;
        }else if (e.compareTo(root.val) > 0) {
            // 右树删除
            root.right = remove(root.right,e);
            retNode = root;
        }else {
            // 当前节点就是待删除节点
            if (root.left == null) {
                // 只有右树，返回右树树根
                Node right = root.right;
                root.right = null;
                size --;
                retNode = right;
            }
            else if (root.right == null) {
                // 只有左树，返回左树树根
                Node left = root.left;
                root.left = null;
                size --;
                retNode = left;
            }
            else {
                // 此时左右树都不为空，找到当前root的后继作为新的树根
                Node successor = minimum(root.right);
                // 在右树中删除该节点
                // 需要注意，此时removeMin这个方法并没有验证平衡性，有可能在这里出现不平衡的树，因此改为remove方法
                successor.right = remove(root.right,successor.val);
                successor.left = root.left;
                root.left = root.right = null;
                // 此处size不用--，因为在removeMin中已经删除了successor，size--了。
                retNode = successor;
            }
        }
        // 有可能删除的是叶子结点，判空
        if (retNode == null) {
            return null;
        }
        // 平衡性处理
        // 更新height
        retNode.height = Math.max(getHeight(retNode.left),getHeight(retNode.right)) + 1;
        // 计算当前节点的平衡因子
        int balancedFactor = getBalancedFactor(retNode);
        // 平衡维护
        // LL
        if (balancedFactor > 1 && getBalancedFactor(retNode.left) >= 0) {
            // 右旋实现，此时插入元素后，不平衡的结点在当前树左侧的左侧,返回旋转后的根节点
            return rightRotate(retNode);
        }

        // RR
        if (balancedFactor < -1 && getBalancedFactor(retNode.right) <= 0) {
            // 左旋实现，此时插入元素后，不平衡的节点在当前树右侧的右侧，返回左旋后根节点
            return leftRotate(retNode);
        }

        // LR - 插入结点后，不平衡的结点在当前结点左树的右侧
        if (balancedFactor > 1 && getBalancedFactor(retNode.left) < 0) {
            // 先将左子树左旋转为LL
            retNode.left = leftRotate(retNode.left);
            // 再将当前root右旋
            return rightRotate(retNode);
        }

        // RL - 插入结点后，不平衡的结点在当前结点右树的左侧
        if (balancedFactor < -1 && getBalancedFactor(retNode.right) > 0) {
            // 先将右子树右旋转为RR
            retNode.right = rightRotate(retNode.right);
            // 再将当前root左旋
            return leftRotate(retNode);
        }
        // 当前结点满足平衡树的定义，继续回溯
        return retNode;
    }

    /**
     * 判断以root为根节点的树中是否包含元素val
     * @param root
     * @param val
     * @return
     */
    private boolean contains(Node root, E val) {
        if (root == null) {
            return false;
        }
        if (root.val.equals(val)) {
            return true;
        }else if (val.compareTo(root.val) < 0) {
            return contains(root.left,val);
        }else {
            return contains(root.right,val);
        }
    }

    public E minimum() {
        if (isEmpty()) {
            throw new NoSuchElementException("BST is empty!");
        }
        return minimum(root).val;
    }

    public E removeMin() {
        E ret = minimum();
        root = removeMin(root);
        return ret;
    }

    /**
     * 删除以root为根节点的BST的最小值节点
     * @param root 当前BST的根
     * @return 删除最小值节点的树根
     */
    private Node removeMin(Node root) {
        if (root.left == null) {
            Node right = root.right;
            root.right = null;
            size --;
            return right;
        }
        root.left = removeMin(root.left);
        return root;
    }

    public E removeMax() {
        E ret = maximum();
        root = removeMax(root);
        return ret;
    }

    /**
     * 删除以root为根节点的BST的最大值节点
     * @param root 当前BST的根
     * @return 删除最大值节点的树根
     */
    private Node removeMax(Node root) {
        if (root.right == null) {
            Node left = root.left;
            root.left = null;
            size --;
            return left;
        }
        root.right = removeMax(root.right);
        return root;
    }

    /**
     * 找到以root为根的二叉树最小值节点
     * @param root
     * @return 最小值节点
     */
    private Node minimum(Node root) {
        if (root.left == null) {
            return root;
        }
        return minimum(root.left);
    }

    public E maximum() {
        if (isEmpty()) {
            throw new NoSuchElementException("BST is empty!");
        }
        return maximum(root).val;
    }

    /**
     * 找到以root为根的二叉树最大值节点
     * @param root
     * @return 最大值节点
     */
    private Node maximum(Node root) {
        if (root.right == null) {
            return root;
        }
        return maximum(root.right);
    }

    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        generateBSTString(root,0,sb);
        return sb.toString();
    }

    /**
     * 先序遍历以当前root为根节点的BST字符串
     * @param root
     * @param depth
     * @param sb
     */
    private void generateBSTString(Node root, int depth, StringBuilder sb) {
        if (root == null) {
            sb.append(generateDepthString(depth)).append("NULL\n");
            return;
        }
        sb.append(generateDepthString(depth)).append(root.val).append("\n");
        generateBSTString(root.left,depth + 1,sb);
        generateBSTString(root.right,depth + 1,sb);
    }

    /**
     * 按深度打印--
     * 每向下一层多两个-
     * @param depth
     * @return
     */
    private String generateDepthString(int depth) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < depth; i++) {
            sb.append("--");
        }
        return sb.toString();
    }
}